package com.hk.commons.crypto;

import com.hk.commons.util.Algorithm;
import com.hk.commons.util.JsonUtils;
import com.hk.commons.util.StringUtils;
import lombok.Getter;
import lombok.SneakyThrows;
import org.apache.commons.codec.binary.Base64;

import javax.crypto.Cipher;
import java.nio.charset.StandardCharsets;
import java.security.*;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Map;
import java.util.Objects;

/**
 * RSA 非对称加密工具
 *
 * @author Kevin
 */
public class RSAHelper {

    public static final String SIGN_ALGORITHMS = "SHA1WithRSA";

    public static final String SHA256_SIGN_ALGORITHMS = "SHA256WithRSA";

    @Getter
    private RSAPrivateKey rsaPrivateKey;

    @Getter
    private RSAPublicKey rsaPublicKey;

    private final Provider provider;

    /**
     * {@link #SIGN_ALGORITHMS } or {@link #SHA256_SIGN_ALGORITHMS}
     */
    private final String algorithm;

    public RSAHelper(String privateKey, String publicKey) {
        this(privateKey, publicKey, SHA256_SIGN_ALGORITHMS);
    }

    public RSAHelper(RSAPrivateKey privateKey, RSAPublicKey publicKey) {
        this(privateKey, publicKey, null, SHA256_SIGN_ALGORITHMS);
    }

    public RSAHelper(String privateKey, String publicKey, String algorithm) {
        this(privateKey, publicKey, null, algorithm);
    }

    public RSAHelper(RSAPrivateKey privateKey, RSAPublicKey publicKey,
                     Provider provider, String algorithm) {
        this.rsaPrivateKey = privateKey;
        this.rsaPublicKey = publicKey;
        this.provider = provider;
        this.algorithm = algorithm;
    }

    public RSAHelper(String privateKey, String publicKey,
                     Provider provider, String algorithm) {
        if (StringUtils.isNotEmpty(privateKey)) {
            this.rsaPrivateKey = getPrivateKey(privateKey);
        }
        if (StringUtils.isNotEmpty(publicKey)) {
            this.rsaPublicKey = getPublicKey(publicKey);
        }
        this.provider = provider;
        this.algorithm = algorithm;
    }

    public Cipher getCipher() {
        return getCipher(this.provider);
    }

    @SneakyThrows(value = {Exception.class})
    public static Cipher getCipher(Provider provider) {
        return Objects.isNull(provider) ? Cipher.getInstance(Algorithm.RSA.getName())
                : Cipher.getInstance(Algorithm.RSA.getName(), provider);
    }

    /**
     * 对内容进行签名
     *
     * @param content 要签名的内容
     */
    public byte[] sign(String content) {
        return sign(content.getBytes(StandardCharsets.UTF_8));
    }

    /**
     * 对内容进行签名
     *
     * @param content 要签名的内容
     */
    @SneakyThrows
    public byte[] sign(byte[] content) {
        var signature = Signature.getInstance(this.algorithm);
        signature.initSign(rsaPrivateKey);
        signature.update(content);
        return signature.sign();
    }

    /**
     * 根据私钥对内容加密
     *
     * @param data 要加密的内容
     */
    public byte[] encryptByPrivateKey(String data) {
        return encrypt(data, this.rsaPrivateKey);
    }

    /**
     * 根据私钥对内容加密,转换成String
     *
     * @param data 要加密的内容
     */
    public String encryptByPrivateKeyString(String data) {
        return new String(encryptByPrivateKey(data));
    }

    /**
     * 根据私钥对内容加密
     *
     * @param data 要加密的内容
     */
    public byte[] encryptByPublicKey(String data) {
        return encrypt(data, this.rsaPublicKey);
    }

    /**
     * 根据私钥对内容加密,转换成String
     *
     * @param data 要加密的内容
     */
    public String encryptByPublicKeyString(String data) {
        return new String(encryptByPublicKey(data));
    }

    /**
     * 对内容加密
     *
     * @param data 要加密的内容
     * @param key  指定加密的 key
     */
    @SneakyThrows
    public byte[] encrypt(String data, Key key) {
        var cipher = getCipher();
        cipher.init(Cipher.ENCRYPT_MODE, key);
        return cipher.doFinal(data.getBytes(StandardCharsets.UTF_8));
    }

    /**
     * 使用公钥对内容解密
     *
     * @param data 要解密的内容
     */
    public byte[] decryptByPublicKey(String data) {
        return decrypt(data, this.rsaPublicKey);
    }

    /**
     * 使用公钥对内容解密
     *
     * @param data 要解密的内容
     */
    public String decryptByPublicKeyString(String data) {
        return new String(decrypt(data, this.rsaPublicKey));
    }

    /**
     * 使用私钥对内容解密
     *
     * @param data 要解密的内容
     */
    public byte[] decryptByPrivateKey(String data) {
        return decrypt(data, this.rsaPrivateKey);
    }

    /**
     * 使用私钥对内容解密，返回String类型
     *
     * @param data 要解密的内容
     */
    public String decryptByPrivateKeyString(String data) {
        return new String(decryptByPrivateKey(data));
    }

    /**
     * 对内容解密
     *
     * @param data 要解密的内容
     * @param key  要解密的 key
     */
    @SneakyThrows
    public byte[] decrypt(String data, Key key) {
        var cipher = getCipher();
        cipher.init(Cipher.DECRYPT_MODE, key);
        return cipher.doFinal(Base64.decodeBase64(data));
    }

    /**
     * 校验签名是否合法，数据是否被修改
     *
     * @param content 内容
     * @param sign    签名
     */
    public boolean verify(String content, String sign) {
        try {
            var signature = Signature.getInstance(this.algorithm);
            signature.initVerify(this.rsaPublicKey);
            signature.update(content.getBytes(StandardCharsets.UTF_8));
            return signature.verify(Base64.decodeBase64(sign));
        } catch (Exception e) {
            return false;
        }
    }

    /**
     * 获取 私钥
     *
     * @param privateKey privateKey
     */
    @SneakyThrows
    public static RSAPrivateKey getPrivateKey(String privateKey) {
        var keyFactory = KeyFactory.getInstance(Algorithm.RSA.getName());
        var keySpec = new PKCS8EncodedKeySpec(Base64.decodeBase64(privateKey));
        return (RSAPrivateKey) keyFactory.generatePrivate(keySpec);
    }

    /**
     * 获取 公钥
     *
     * @param publicKey publicKey
     */
    @SneakyThrows
    public static RSAPublicKey getPublicKey(String publicKey) {
        var keyFactory = KeyFactory.getInstance(Algorithm.RSA.getName());
        var keySpec = new X509EncodedKeySpec(Base64.decodeBase64(publicKey));
        return (RSAPublicKey) keyFactory.generatePublic(keySpec);
    }

    /**
     * 测试: 生成公钥与私钥
     *
     * @param keySize 1024 or 2048
     */
    @SneakyThrows
    public static Map<String, String> makeKeyPair(int keySize) {
        var keyPairGenerator = KeyPairGenerator.getInstance(Algorithm.RSA.getName());
        keyPairGenerator.initialize(keySize);
        KeyPair keyPair = keyPairGenerator.generateKeyPair();
        var privateKey = (RSAPrivateKey) keyPair.getPrivate();
        var publicKey = (RSAPublicKey) keyPair.getPublic();
        return Map.of("privateKey", Base64.encodeBase64String(privateKey.getEncoded()),
                "publicKey", Base64.encodeBase64String(publicKey.getEncoded()));
    }

    /**
     * 测试: 生成公钥与私钥
     *
     * @param keySize 1024 or 2048 or 4096
     */
    @SneakyThrows
    public static KeyPair keyPair(int keySize) {
        var keyPairGenerator = KeyPairGenerator.getInstance(Algorithm.RSA.getName());
        keyPairGenerator.initialize(keySize);
        var keyPair = keyPairGenerator.generateKeyPair();
        var privateKey = (RSAPrivateKey) keyPair.getPrivate();
        var publicKey = (RSAPublicKey) keyPair.getPublic();
        return new KeyPair(publicKey, privateKey);
    }

    public static void main(String[] args) {
        System.out.println(JsonUtils.serialize(makeKeyPair(2048), true));
    }

}